import argparse
from pathlib import Path
import re
import sys
sys.path.append("/home/lxy/DPR")
import retriever
from xinwangban_fuzzy_re import fuzzy_search
from functools import partial
import json
# index_dir={'neu':'/home/lxy/DPR/9m19d_neu_1w5_bge1.5','gaobao':'/home/lxy/DPR/9m9d_gaobao_18w_norm'}

index_dir='/data/lxy/RAT/10m9d_wxb_bge1.5_hnswivf_20500'
# model_name={'neu':'BAAI/bge-large-zh-v1.5','gaobao':'/home/lxy/DPR/models/checkpoint-53910-epoch-10'}
model_name='BAAI/bge-large-zh-v1.5'
r=retriever.retriver(index_dir,model_name,only_retriever=True)

with open('/home/lxy/wxbdata/final_test.jsonl','r',encoding='utf-8') as f:
    data=list(map(partial(json.loads),f.readlines()))

tp,fp,tn,fn=0.000001,0.000001,0.000001,0.000001
correct=0
jiexi,wangluo=0,0
for i,d in enumerate(data):
    input=d['question']
    target=d['name']
    # result=r.retrieve(input,k=1,debug=False)
    # if result == 0:
    try:
        result,_=fuzzy_search(input)
        if len(result):
            result=result[0]
    except Exception as e:
        result=""
        wangluo+=1
        print(e)

    match=None
    try:
        match = re.search(r'业务名称：(.+)', result)
        if match:
            match=match.group(1)
            # print(i,match)
        else:
            # print('解析失败',i,match,result)
            jiexi+=1
            match=None
    except:
        # print('解析失败',i,match,result)
        jiexi+=1



    if target !="None":
        if match==target:
            tp+=1
            correct+=1
        else:
            fp+=1
    else:
        if match is None:
            print(i)
            tn+=1
            correct+=1
        else:
            fn+=1
    
print(tp,tn,fp,fn)
precision=tp/(tp+fp)
recall=tp/(tp+fn)
print('acc:',correct/len(data))
print('precision:',precision)
print('recall',recall)
print('f1:',2*precision*recall/(precision+recall))
print('解析失败',jiexi/len(data))
print('网络失败',wangluo/len(data))